/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.io.*;
import java.text.DecimalFormat;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.*;

import edu.umd.clip.lm.factors.FactorTuple;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class TreeTrellisNode {
	private final Lock lock = new ReentrantLock();
	final long currentWord;
	private double totalProb;
	private double logScale;
	private final double prevLogScale;
	final int time;
	private boolean decoded = false;
	private boolean deleted = false;
	
	private final TreeTrellisNode previousNode;
	private final static Lock nodeLinkLock = new ReentrantLock();
	private final AtomicInteger useCount = new AtomicInteger(1);
	private final TrellisNodeCache cache;
	
	private Map<FactorTuple, TreeTrellisNode> nextNodes;
	IdentityHashMap<TrellisNodeItem, TrellisNodeItem> items;

	public TreeTrellisNode() {
		this(null);
	}
	
	public TreeTrellisNode(TrellisNodeCache cache) {
		currentWord = 0;
		totalProb = 1.0;
		logScale = 0.0;
		prevLogScale = 0.0;
		time = 0;
		previousNode = null;
		nextNodes = new HashMap<FactorTuple, TreeTrellisNode>();
		items = null;
		this.cache = cache; 
	}
	
	/**
	 * @param trellis
	 * @param previousNode
	 * @param currentWord
	 */
	TreeTrellisNode(TreeTrellisNode previousNode, long currentWord) {
		this.currentWord = currentWord;
		this.previousNode = previousNode;
		this.totalProb = Double.NaN;
		this.logScale = 0.0;
		this.prevLogScale = previousNode == null ? 0.0 : previousNode.getCumulativeLogScale();
		this.time = previousNode == null ? 1 : previousNode.time + 1;
		
		this.nextNodes = new HashMap<FactorTuple, TreeTrellisNode>();
		this.items = new IdentityHashMap<TrellisNodeItem, TrellisNodeItem>();
		this.cache = previousNode.cache;
	}

	@SuppressWarnings("unchecked")
	protected void split(TrellisNodeItem item, TrellisNodeItem cutOff) {
		// copy links from the previous slice
		IdentityHashMap<TrellisNodeItem,TrellisNodeItem> itemPrevLinks = item.getPrevItemsMap();
		if (itemPrevLinks != null) {
			cutOff.setPrevItems((IdentityHashMap<TrellisNodeItem,TrellisNodeItem>) itemPrevLinks.clone());
			for(TrellisNodeItem i : item.getPrevItems()) {
				i.addNextItem(cutOff);
			}
		}
		// copy links to the following slice if it exists
		IdentityHashMap<TrellisNodeItem,TrellisNodeItem> itemNextLinks = item.getNextItemsMap();
		if (itemNextLinks != null) {
			cutOff.setNextItems((IdentityHashMap<TrellisNodeItem,TrellisNodeItem>) itemNextLinks.clone());
			for(TrellisNodeItem i : cutOff.getNextItems()) {
				i.addPrevItem(cutOff);
			}
		}
		items.put(cutOff, cutOff);
	}

	protected void insertItem(TrellisNodeItem item) {
		if (time != item.time) {
			System.out.print("");
		}
		items.put(item, item);
		if (!isDecoded()) {
			// BUG: if the node is already decoded and computeTotalProb() is *not* invoked right after 
			// the item is inserted and probabilities are adjusted,
			// a later invocation of computeTotalProb() returns wrong results for some reason
			this.totalProb = Double.NaN;
		}
	}
	
	protected void removeItem(TrellisNodeItem item) {
		double prob = item.getProb();
		this.totalProb -= prob;
		
		items.remove(item);
		
		if (item.getPrevItems() != null) {
			for(TrellisNodeItem prevItem : item.getPrevItems()) {
				prevItem.removeNextItem(item);
			}
		}
		if (item.getNextItems() != null) {
			for(TrellisNodeItem nextItem : item.getNextItems()) {
				nextItem.removeNextItem(item);
			}
		}
	}
	
	protected void lock() {
		//System.err.printf("%s locking node %s\n", Thread.currentThread().getName(), toString());
		//System.err.flush();
		lock.lock();
		//System.err.printf("%s locked node %s\n", Thread.currentThread().getName(), toString());
		//System.err.flush();
	}
	
	protected boolean tryLock() {
		//System.err.printf("%s trying to lock node %s\n", Thread.currentThread().getName(), toString());
		boolean locked = lock.tryLock();
		//System.err.printf("%s %slocked node %s\n", Thread.currentThread().getName(), locked? "" : "NOT ", toString());
		//System.err.flush();
		return locked;
	}
	
	protected void unlock() {
		//System.err.printf("%s unlocked node %s\n", Thread.currentThread().getName(), toString());
		//System.err.flush();
		lock.unlock();
	}

	protected void lock(int depth) {
		/*
		TreeTrellisNode node = this;
		while(!node.isHead() && depth > 0) {
			node.lock();
			--depth;
			node = node.getPreviousNode();
		}
		*/
		if (isHead() || depth <= 0) return;
		lock();
		previousNode.lock(depth-1);
	}
	
	protected boolean tryLock(int depth) {
		if (depth <= 0 || isHead()) return true;
		if (tryLock()) {
			boolean result = previousNode.tryLock(depth-1);
			if (!result) unlock();
			return result;
		}
		return false;
	}
	
	protected void unlock(int depth) {
		TreeTrellisNode node = this;
		while(!node.isHead() && depth > 0) {
			node.unlock();
			--depth;
			node = node.getPreviousNode();
		}
	}
	
	private void dropLocks(int depth, int num) {
		TreeTrellisNode node = this;
		while(node != null && depth > 0) {
			--depth;
			node = node.getPreviousNode();
		}
		if (node != null) {
			node.unlock(num);
		}
	}
	
	private double computeTotalProb() {
		double p = 0;
		for(TrellisNodeItem item : items.keySet()) {
			p += item.getProb();
		}
		return p;
	}
	
	private void printItems() {
		double total = 0;
		for(TrellisNodeItem item : items.keySet()) {
			double prevProb = item.getPrevProb();
			double itemProb = item.getProb() / prevProb;
			System.err.printf("%s: prev=%g, item=%g\n", item, prevProb, itemProb);
			total += item.getProb();
		}		
		System.err.printf("Total prob: %g, scale: %g\n", total, logScale);
	}
	
	double getTotalProb() {
		if (Double.isNaN(totalProb)) {
			totalProb = computeTotalProb();
		}
		return totalProb;
	}

	void setTotalProb(double totalProb) {
		this.totalProb = totalProb;
	}

	double getLogScale() {
		return logScale;
	}

	void setLogScale(double logScale) {
		this.logScale = logScale;
	}

	double getLogProb() {
		return getCumulativeLogScale() + Math.log10(getTotalProb());
	}
	
	double getCumulativeLogScale() {
		return prevLogScale + logScale;
	}

	TreeTrellisNode getPreviousNode() {
		return previousNode;
	}

	long getCurrentWord() {
		return currentWord;
	}
	
	TrellisNodeItem[] itemsArray() {
		return items.keySet().toArray(new TrellisNodeItem[items.size()]);
	}
	
	protected void rescale() {
		double maxScale = 0.0;
		double totalProb = 0.0;
		for(TrellisNodeItem item : items.keySet()) {
			double prob = item.getProb();
			totalProb += prob;
			double scale = item.tree.getProbTree().getScale();
			if (scale > maxScale) {
				maxScale = scale;
			}
		}
		double scale = maxScale > 0.0 ? 1.0 / maxScale : 1.0;

		for(TrellisNodeItem item : items.keySet()) {
			item.tree.scale(scale);
		}
		this.totalProb = totalProb * scale;
		this.logScale -= Math.log10(scale);
	}
	
	public TreeTrellisNode appendNodeAndLock(long word) {
		// minimum depth to lock the current node
		//int tmpDepth = depth > 1 ? depth : 2;
		
		FactorTuple tuple = new FactorTuple(word);

		useCount.incrementAndGet();
		/*
		if (cache != null) {
			cache.touch(this);
		}
		*/
		TreeTrellisNode theNode;
		
		boolean needToLock = false;
		nodeLinkLock.lock();
		try {
			if (cache != null) {
				cache.addRequest(time + 1);
			}
			
			TreeTrellisNode oldNode = nextNodes.get(tuple);
			if (oldNode != null) {
				// don't try to acquire a node lock while holding nodeLinkLock
				//oldNode.lock();
				needToLock = true;
				theNode = oldNode;
				if (cache != null) {
					cache.addHit(time + 1);
				}
			} else {
				TreeTrellisNode node = new TreeTrellisNode(this, word);
				node.lock();
				nextNodes.put(tuple, node);
				theNode = node;
			}
		} finally {
			nodeLinkLock.unlock();
		}
		/*
		if (cache != null) {
			cache.touch(theNode);
		}
		*/
		if (needToLock) theNode.lock();
		return theNode;
	}

	public void dumpNodePath(String filename, boolean includeDeadEnd) {
		try {
			dumpNodePath(new FileOutputStream(filename), includeDeadEnd);
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		}
	}
	
	public void dumpNodePath(OutputStream os, boolean includeDeadEnd) {
		OutputStreamWriter out = new OutputStreamWriter(os);
		DecimalFormat nf = new DecimalFormat("0.###E0");
		nf.setMaximumFractionDigits(3);

		ArrayDeque<TreeTrellisNode> nodes = new ArrayDeque<TreeTrellisNode>(time + 2);
		for(TreeTrellisNode node = this; node.getPreviousNode() != null; node = node.getPreviousNode()) {
			nodes.addFirst(node);
		}
		
		HashSet<TrellisNodeItem> deadEnds = new HashSet<TrellisNodeItem>();
		try {
			out.append("digraph Trellis {\n");
			out.append("  rankdir=LR\n;");
			
			for(TreeTrellisNode node : nodes) {
				out.append(String.format("// time=%d, word=%s\n", node.time, FactorTuple.toStringNoNull(node.currentWord)));
				deadEnds.clear();
				for(TrellisNodeItem item : node.items.keySet()) {
					if (!includeDeadEnd && node != this) {
						if (item.getNextItems() == null || item.getNextItems().size() == 0) {
							// drop items which have to continuation in the next slice
							deadEnds.add(item);
							continue;
						}
					}
					out.append("  N" + item.toString());
					out.append(" [label=\"N" + item.toString());
					/*
					if (!item.clusterid.isEmpty()) {
						out.append(" (" + item.clusterid + ")");
					}
					*/
					try {
						String prefix = item.tree.getPrefix().toString();
						if (!prefix.isEmpty()) {
							out.append("\\nprefix=");
							out.append(prefix);
						}
					} catch(UnsupportedOperationException e) {
						// item.tree.getPrefix() will throw this if it's a ghost
					}
					
					out.append("\\nprev=" + nf.format(item.getPrevProb() * Math.pow(10.0, node.getCumulativeLogScale() - node.getLogScale())));
					out.append("\\nprob=" + nf.format(item.tree.getTotalProb() * Math.pow(10.0, node.getLogScale())) + "\"];\n");
				}
				out.append(" {rank=same;");
				for(TrellisNodeItem item : node.items.keySet()) {
					if (includeDeadEnd || !deadEnds.contains(item)) {
						out.append("  N" + item.toString());
					}
				}
				out.append("};\n");
				
				for(TrellisNodeItem item : node.items.keySet()) {
					if (!includeDeadEnd && deadEnds.contains(item)) {
						continue;
					}
					Set<TrellisNodeItem> prevLinks = item.getPrevItems();
					if (prevLinks != null) {
						out.append("N" + item.toString() + " -> {");
						for(TrellisNodeItem prev : prevLinks) {
							out.append(" N" + prev.toString());
						}
						out.append("};\n");
					}
					Set<TrellisNodeItem> nextLinks = item.getNextItems();
					if (nextLinks != null) {
						out.append("N" + item.toString() + " -> {");
						for(TrellisNodeItem next : nextLinks) {
							out.append(" N" + next.toString());
						}
						out.append("};\n");
					}
				}
			}
			out.append("}\n");
			out.flush();

		} catch(IOException e) {
			
		}
	}
	
	public String toString() {
		return String.format("%s@%d", FactorTuple.toStringNoNull(currentWord), time);
	}
	
	public final boolean isHead() {
		return previousNode == null;
	}

	public final boolean isDecoded() {
		return decoded;
	}

	public final boolean isDeleted() {
		return deleted;
	}
	
	public final void finishDecoding() {
		this.decoded = true;
	}
	
	public void releaseNode() {
		if (useCount.decrementAndGet() == 0 && cache != null) {
			cache.touch(this);
		}
	}
	
	public void releasePath() {
		TreeTrellisNode node = this;
		while(node != null) {
			node.releaseNode();
			node = node.getPreviousNode();
		}
	}
	
	public boolean pruneNode() {
		if (!(decoded && nextNodes.isEmpty() && useCount.get() == 0 && !isHead())) {
			return false;
		}
		// delete the node
		if (tryLock(2)) {
			try {
				nodeLinkLock.lock();
				try {
					if (!(decoded && nextNodes.isEmpty() && useCount.get() == 0 && !isHead())) {
						return false;
					}
					previousNode.nextNodes.remove(new FactorTuple(currentWord));
					previousNode.useCount.decrementAndGet();
					deleted = true;
				} finally {
					nodeLinkLock.unlock();
				}
				for(TrellisNodeItem item : items.keySet()) {
					Set<TrellisNodeItem> prevItems = item.getPrevItems();
					if (prevItems != null) {
						for(TrellisNodeItem prevItem : prevItems) {
							prevItem.removeNextItem(item);
						}
					}
				}
				if (cache != null) {
					cache.addPruned(time);
				}
			} finally {
				unlock(2);
			}
			return true;
		} 
		return false;
	}
	
	public int getUseCount() {
		return useCount.get();
	}

	public final TrellisNodeCache getCache() {
		return cache;
	}
}
